
rm(list = ls())

## setwd("/Users/lichengl/Desktop/BSA_replication_code/")


library(modeest)

load("results/sim_N50_T30.RData")
load("results/sim_N50_T50.RData")
load("results/sim_N50_T90.RData")

load("results/sim_N100_T30.RData")
load("results/sim_N100_T50.RData")
load("results/sim_N100_T90.RData")




## coverage rate and average length of 95% CI 
cover <- function(x) {
    nboots <- dim(x)[1]
    q <- sapply(1:nboots, function(vec) {x[vec, 1] <= 0 & x[vec, 2] >= 0} )
    return(mean(q)) ##  * 100)
}

## length of 95% CI
len <- function(x) {
    nboots <- dim(x)[1]
    q <- sapply(1:nboots, function(vec) {x[vec, 2] - x[vec, 1]} )
    return(mean(q))
}




## 4 models and 6 combinations of unit number and time period 
model <- rep(c("Naive", "BSA w/ Shrinkage", "BSA w/o Shrinkage", "Oracle"), 6)
## number of units
Nu <- rep(c(50, 100), each = 12)
## number of time periods
TT <- rep(rep(c(30, 50, 90), each = 4), 2)

## bias, rmse, standard error
bias <- rmse <- se <- c()

## simulation results
mc_all <- list(mc_50_30, mc_50_50, mc_50_90,
	           mc_100_30, mc_100_50, mc_100_90)



mc_list_95 <- list()
mc_list_50 <- list()

## generate cis 
for (k in 1:6) {
	mc_sub <- mc_all[[k]]
	for (j in 1:4) {
		mc_sub_m <- mc_sub[j,,] ## different models ## 10000 * 1000 
		sub_ci <- apply(mc_sub_m, 2, quantile, c(0.025, 0.975)) ## 2 * 1000
		sub_ci2 <- apply(mc_sub_m, 2, quantile, c(0.25, 0.75)) ## 2 * 1000
		
		mc_list_95 <- c(mc_list_95, list(t(sub_ci)))
		mc_list_50 <- c(mc_list_50, list(t(sub_ci2)))


	}
}

cover_95 <- sapply(1:24, function(vec) {return(cover(mc_list_95[[vec]]))})
length_95 <- sapply(1:24, function(vec) {return(len(mc_list_95[[vec]]))})

cover_50 <- sapply(1:24, function(vec) {return(cover(mc_list_50[[vec]]))})
length_50 <- sapply(1:24, function(vec) {return(len(mc_list_50[[vec]]))})




data <- cbind.data.frame(Nu, TT, model, cover_95, length_95, cover_50, length_50)
names(data) <- c("N", "T", "model", "cover95", "len95", "cover50", "len50")


## CI length
data[, c(5, 7)] <- lapply(data[, c(5, 7)], round, 3)


## -------------------------------------------------- ##
##               1. Table replication                 ##
## -------------------------------------------------- ##  

## data 

#     N  T             model cover95 len95 cover50 len50
#1   50 30             Naive    57.7 1.149    15.8 0.395
#2   50 30  BSA w/ Shrinkage   100.0 4.350    78.9 1.010
#3   50 30 BSA w/o Shrinkage   100.0 5.379    92.4 1.423
#4   50 30            Oracle    91.5 1.129    43.6 0.388
#5   50 50             Naive    52.6 1.043    11.8 0.359
#6   50 50  BSA w/ Shrinkage   100.0 4.336    80.8 0.988
#7   50 50 BSA w/o Shrinkage   100.0 5.381    94.8 1.417
#8   50 50            Oracle    93.9 1.022    49.3 0.352
#9   50 90             Naive    47.2 0.985     9.1 0.339
#10  50 90  BSA w/ Shrinkage   100.0 4.340    81.8 0.977
#11  50 90 BSA w/o Shrinkage   100.0 5.384    95.7 1.414
#12  50 90            Oracle    93.6 0.961    48.5 0.331
#13 100 30             Naive    34.6 0.804     4.5 0.277
#14 100 30  BSA w/ Shrinkage   100.0 4.300    87.3 0.941
#15 100 30 BSA w/o Shrinkage   100.0 5.347    98.1 1.387
#16 100 30            Oracle    91.6 0.791    45.7 0.272
#17 100 50             Naive    23.4 0.731     2.2 0.252
#18 100 50  BSA w/ Shrinkage   100.0 4.306    89.8 0.933
#19 100 50 BSA w/o Shrinkage   100.0 5.354    99.1 1.389
#20 100 50            Oracle    93.1 0.717    43.6 0.247
#21 100 90             Naive    16.3 0.690     1.4 0.237
#22 100 90  BSA w/ Shrinkage   100.0 4.321    90.6 0.933
#23 100 90 BSA w/o Shrinkage   100.0 5.360    99.6 1.392
#24 100 90            Oracle    94.8 0.674    47.6 0.232

## table 2, 3, 4, 5
## column 4: table 2
## column 5: table 3
## column 6: table 4
## column 7: table 5

write.csv(data, file = "table/tab2_5.csv")


## power (given different att magnitude)
a <- seq(3, -3.5, length.out = 20)
powers <- function(x, ATT = a) {
    nboots <- dim(x)[1]
    n2 <- length(ATT)

    q <- c()

    for (i in 1:n2) {
    	x0 <- x 
    	x0 <- x0 + ATT[i]
    	#if (ATT[i] > 0) {
    	#	subq <- sapply(1:nboots, function(vec) {x0[vec, 1] > 0} )
    	#} else {
    	#	subq <- sapply(1:nboots, function(vec) {x0[vec, 2] < 0} )
    	#}

    	subq <- sapply(1:nboots, function(vec) {x0[vec, 1] > 0 | x0[vec, 2] < 0} )
    	q <- c(q, mean(subq) * 100)

    }

    return(q)
}


## a matrix of 20 (number of eff) * 24 (number of models 4 and N, T combs 6)
power_all <- sapply(1:24, function(k) {return(powers(mc_list_95[[k]], ATT = a))})

la <- length(a)




## -------------------------------------------------------- ##
##               2. power curve replication                 ##
## -------------------------------------------------------- ##  




## 2.1 N = 50, T = 30 

k <- 1  

##pdf("plot/sim1_power_N50_T30.pdf", width = 10, height = 6)

pdf("plot/figure1_l1.pdf", width = 10, height = 6)
plot(a, 
	 power_all[, k + 1], type = "b", 
	 ylim = c(0, 100), xlim = c(-3.6, 3), col = "red", pch = 19,
	 ylab = "", xlab = "Treatment Effect",
	 cex.axis = 1.5, cex.lab = 1.5)

lines(a, power_all[, k + 2],
	  col = "red", type = "b", lty = 2, pch = 18)
lines(a, power_all[, k + 3],
	  col = "grey", type = "b", lty = 1, pch = 19)
lines(a, power_all[, k],
	  col = "grey", type = "b", lty = 2, pch = 18)

title(ylab = "% of CIs Excluding Zero", 
	  cex.lab = 1.5, line = 2.5)

legend(-3.8, 23, 
	   legend = c("DM-LFM w/ U", "DM-LFM w/o U", "BSA w/ Shrinkage", "BSA w/o Shrinkage"),
       col=c("grey", "grey", "red", "red"), 
       lty = c(1, 2, 1, 2), pch = c(19, 18, 19, 18), cex = rep(1, 4), bty="n")

dev.off()







## 2.2 N = 50, T = 50 

k <- 5  

##pdf("plot/sim1_power_N50_T50.pdf", width = 10, height = 6)
pdf("plot/figure1_l2.pdf", width = 10, height = 6)
plot(a, 
	 power_all[, k + 1], type = "b", 
	 ylim = c(0, 100), xlim = c(-3.6, 3), col = "red", pch = 19,
	 ylab = "", xlab = "Treatment Effect",
	 cex.axis = 1.5, cex.lab = 1.5)

lines(a, power_all[, k + 2],
	  col = "red", type = "b", lty = 2, pch = 18)
lines(a, power_all[, k + 3],
	  col = "grey", type = "b", lty = 1, pch = 19)
lines(a, power_all[, k],
	  col = "grey", type = "b", lty = 2, pch = 18)

title(ylab = "% of CIs Excluding Zero", 
	  cex.lab = 1.5, line = 2.5)

legend(-3.8, 23, 
	   legend = c("DM-LFM w/ U", "DM-LFM w/o U", "BSA w/ Shrinkage", "BSA w/o Shrinkage"),
       col=c("grey", "grey", "red", "red"), 
       lty = c(1, 2, 1, 2), pch = c(19, 18, 19, 18), cex = rep(1, 4), bty="n")

dev.off()







## 2.3 N = 50, T = 90

k <- 9  

##pdf("plot/sim1_power_N50_T90.pdf", width = 10, height = 6)
pdf("plot/figure1_l3.pdf", width = 10, height = 6)
plot(a, 
	 power_all[, k + 1], type = "b", 
	 ylim = c(0, 100), xlim = c(-3.6, 3), col = "red", pch = 19,
	 ylab = "", xlab = "Treatment Effect",
	 cex.axis = 1.5, cex.lab = 1.5)

lines(a, power_all[, k + 2],
	  col = "red", type = "b", lty = 2, pch = 18)
lines(a, power_all[, k + 3],
	  col = "grey", type = "b", lty = 1, pch = 19)
lines(a, power_all[, k],
	  col = "grey", type = "b", lty = 2, pch = 18)

title(ylab = "% of CIs Excluding Zero", 
	  cex.lab = 1.5, line = 2.5)

legend(-3.8, 23, 
	   legend = c("DM-LFM w/ U", "DM-LFM w/o U", "BSA w/ Shrinkage", "BSA w/o Shrinkage"),
       col=c("grey", "grey", "red", "red"), 
       lty = c(1, 2, 1, 2), pch = c(19, 18, 19, 18), cex = rep(1, 4), bty="n")

dev.off()








## 2.4 N = 100, T = 30

k <- 13  

##pdf("plot/sim1_power_N100_T30.pdf", width = 10, height = 6)
pdf("plot/figur1_r1.pdf", width = 10, height = 6)
plot(a, 
	 power_all[, k + 1], type = "b", 
	 ylim = c(0, 100), xlim = c(-3.6, 3), col = "red", pch = 19,
	 ylab = "", xlab = "Treatment Effect",
	 cex.axis = 1.5, cex.lab = 1.5)

lines(a, power_all[, k + 2],
	  col = "red", type = "b", lty = 2, pch = 18)
lines(a, power_all[, k + 3],
	  col = "grey", type = "b", lty = 1, pch = 19)
lines(a, power_all[, k],
	  col = "grey", type = "b", lty = 2, pch = 18)

title(ylab = "% of CIs Excluding Zero", 
	  cex.lab = 1.5, line = 2.5)

legend(-3.8, 23, 
	   legend = c("DM-LFM w/ U", "DM-LFM w/o U", "BSA w/ Shrinkage", "BSA w/o Shrinkage"),
       col=c("grey", "grey", "red", "red"), 
       lty = c(1, 2, 1, 2), pch = c(19, 18, 19, 18), cex = rep(1, 4), bty="n")

dev.off()








## 2.5 N = 100, T = 50

k <- 17  

##pdf("plot/sim1_power_N100_T50.pdf", width = 10, height = 6)
pdf("plot/figure1_r2.pdf", width = 10, height = 6)
plot(a, 
	 power_all[, k + 1], type = "b", 
	 ylim = c(0, 100), xlim = c(-3.6, 3), col = "red", pch = 19,
	 ylab = "", xlab = "Treatment Effect",
	 cex.axis = 1.5, cex.lab = 1.5)

lines(a, power_all[, k + 2],
	  col = "red", type = "b", lty = 2, pch = 18)
lines(a, power_all[, k + 3],
	  col = "grey", type = "b", lty = 1, pch = 19)
lines(a, power_all[, k],
	  col = "grey", type = "b", lty = 2, pch = 18)

title(ylab = "% of CIs Excluding Zero", 
	  cex.lab = 1.5, line = 2.5)

legend(-3.8, 23, 
	   legend = c("DM-LFM w/ U", "DM-LFM w/o U", "BSA w/ Shrinkage", "BSA w/o Shrinkage"),
       col=c("grey", "grey", "red", "red"), 
       lty = c(1, 2, 1, 2), pch = c(19, 18, 19, 18), cex = rep(1, 4), bty="n")

dev.off()







## 2.6 N = 100, T = 90

k <- 21  

##pdf("plot/sim1_power_N100_T90.pdf", width = 10, height = 6)
pdf("plot/figure1_r3.pdf", width = 10, height = 6)
plot(a, 
	 power_all[, k + 1], type = "b", 
	 ylim = c(0, 100), xlim = c(-3.6, 3), col = "red", pch = 19,
	 ylab = "", xlab = "Treatment Effect",
	 cex.axis = 1.5, cex.lab = 1.5)

lines(a, power_all[, k + 2],
	  col = "red", type = "b", lty = 2, pch = 18)
lines(a, power_all[, k + 3],
	  col = "grey", type = "b", lty = 1, pch = 19)
lines(a, power_all[, k],
	  col = "grey", type = "b", lty = 2, pch = 18)

title(ylab = "% of CIs Excluding Zero", 
	  cex.lab = 1.5, line = 2.5)

legend(-3.8, 23, 
	   legend = c("DM-LFM w/ U", "DM-LFM w/o U", "BSA w/ Shrinkage", "BSA w/o Shrinkage"),
       col=c("grey", "grey", "red", "red"), 
       lty = c(1, 2, 1, 2), pch = c(19, 18, 19, 18), cex = rep(1, 4), bty="n")

dev.off()


